package org.example;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;

public class NIOClient {
    private Selector selector;
    private SocketChannel socketChannel;
    private volatile boolean stop;

    public NIOClient(String host, int port) throws IOException {
        selector = Selector.open();
        socketChannel = SocketChannel.open();
        socketChannel.configureBlocking(false);
        socketChannel.connect(new InetSocketAddress(host, port));
        socketChannel.register(selector, SelectionKey.OP_CONNECT | SelectionKey.OP_READ);
    }

    public void start() throws IOException {
        while (!stop) {
            selector.select();
            Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
            while (iterator.hasNext()) {
                SelectionKey key = iterator.next();
                iterator.remove();

                if (!key.isValid()) continue;

                if (key.isConnectable()) {
                    if (socketChannel.isConnectionPending()) {
                        if (socketChannel.finishConnect()) {
                            socketChannel.register(selector, SelectionKey.OP_READ);
                            doWrite(socketChannel);
                        } else {
                            System.out.println("Failed to connect.");
                            stop = true;
                        }
                    }
                }

                if (key.isReadable()) {
                    ByteBuffer buffer = ByteBuffer.allocate(1024);
                    int numRead = socketChannel.read(buffer);
                    if (numRead > 0) {
                        buffer.flip();
                        byte[] bytes = new byte[buffer.remaining()];
                        buffer.get(bytes);
                        String message = new String(bytes, StandardCharsets.UTF_8);
                        System.out.println("Received from server: " + message);
                        if (message.equals("Bye!")) {
                            stop = true;
                        }
                    } else if (numRead < 0) {
                        key.cancel();
                        socketChannel.close();
                    }
                }
            }
        }

        if (selector != null) {
            selector.close();
        }
    }

    private void doWrite(SocketChannel channel) throws IOException {
        byte[] request = "Hello, Server!".getBytes(StandardCharsets.UTF_8);
        ByteBuffer buffer = ByteBuffer.allocate(request.length);
        buffer.put(request);
        buffer.flip();
        channel.write(buffer);
    }

    public static void main(String[] args) {
        try {
            NIOClient client = new NIOClient("localhost", 8080);
            client.start();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}